package com.mydemo.utils.base2.vosk_model;

import org.vosk.Model;
import org.vosk.Recognizer;
import javax.sound.sampled.*;
import java.io.*;
import java.util.*;
import java.util.concurrent.*;

public class MultiModelTranscriber {
    private List<Model> models;
    private List<String> modelPaths;
    private ExecutorService executorService;
    private String outputTextPath;

    public MultiModelTranscriber() {
        // 初始化模型路径列表
        modelPaths = Arrays.asList(
                "D:\\mp4totext\\vosk-model\\vosk-model-cn-0.22",
                "D:\\mp4totext\\vosk-model\\vosk-model-cn-kaldi-multicn-0.15",
                "D:\\mp4totext\\vosk-model\\vosk-model-small-cn-0.22"
        );
        models = new ArrayList<>();
        executorService = Executors.newFixedThreadPool(modelPaths.size());

        // 设置默认输出路径
        this.outputTextPath = "D:\\mp4totext\\transcription_results.txt";
        loadModels();
    }

    /**
     * 设置输出文本文件路径
     */
    public void setOutputTextPath(String outputTextPath) {
        this.outputTextPath = outputTextPath;
    }

    /**
     * 加载所有语音识别模型
     */
    private void loadModels() {
        System.out.println("开始加载语音识别模型...");

        for (String modelPath : modelPaths) {
            try {
                File modelDir = new File(modelPath);
                if (!modelDir.exists()) {
                    System.err.println("模型路径不存在: " + modelPath);
                    continue;
                }

                System.out.println("加载模型: " + modelPath);
                Model model = new Model(modelPath);
                models.add(model);
                System.out.println("✓ 模型加载成功: " + modelPath);

            } catch (Exception e) {
                System.err.println("✗ 模型加载失败: " + modelPath + " - " + e.getMessage());
                e.printStackTrace();
            }
        }

        if (models.isEmpty()) {
            throw new RuntimeException("没有成功加载任何模型！");
        }

        System.out.println("模型加载完成，成功加载 " + models.size() + " 个模型");
    }

    /**
     * 使用多模型融合识别音频文件
     */
    public String transcribeWithConsensus(String audioPath) {
        if (models.isEmpty()) {
            throw new IllegalStateException("没有可用的模型");
        }

        System.out.println("\n开始多模型融合识别: " + audioPath);

        try {
            // 详细检查音频文件
            AudioFileInfo audioInfo = analyzeAudioFile(audioPath);
            System.out.println("音频文件分析结果:");
            System.out.println("  文件大小: " + audioInfo.fileSize + " 字节");
            System.out.println("  时长: " + audioInfo.duration + " 秒");
            System.out.println("  采样率: " + audioInfo.sampleRate + " Hz");
            System.out.println("  声道数: " + audioInfo.channels);
            System.out.println("  编码格式: " + audioInfo.encoding);

            if (audioInfo.fileSize == 0) {
                throw new RuntimeException("音频文件为空或损坏");
            }

            // 使用多线程并行识别
            final String finalAudioPath = audioPath;
            List<Future<TranscriptionResult>> futures = new ArrayList<>();
            for (Model model : models) {
                Future<TranscriptionResult> future = executorService.submit(() -> {
                    return transcribeWithModel(model, finalAudioPath);
                });
                futures.add(future);
            }

            // 收集所有识别结果
            List<String> results = new ArrayList<>();
            List<TranscriptionResult> transcriptionResults = new ArrayList<>();

            for (int i = 0; i < futures.size(); i++) {
                try {
                    Future<TranscriptionResult> future = futures.get(i);
                    TranscriptionResult result = future.get(5, TimeUnit.MINUTES);
                    transcriptionResults.add(result);
                    results.add(result.text);

                    System.out.println("\n模型 " + (i + 1) + " 识别结果:");
                    System.out.println("  文本: '" + result.text + "'");
                    System.out.println("  文本长度: " + result.text.length());
                    System.out.println("  置信度: " + result.confidence);
                    System.out.println("  处理状态: " + result.status);

                    // 显示所有识别片段
                    if (!result.segments.isEmpty()) {
                        System.out.println("  识别片段 (" + result.segments.size() + "个):");
                        for (int j = 0; j < result.segments.size(); j++) {
                            System.out.println("    " + (j+1) + ". '" + result.segments.get(j) + "'");
                        }
                    }

                } catch (TimeoutException e) {
                    System.err.println("模型 " + (i + 1) + " 识别超时");
                    results.add("");
                    transcriptionResults.add(new TranscriptionResult("", 0.0, "模型" + (i+1), "超时", new ArrayList<>()));
                } catch (Exception e) {
                    System.err.println("模型 " + (i + 1) + " 识别失败: " + e.getMessage());
                    results.add("");
                    transcriptionResults.add(new TranscriptionResult("", 0.0, "模型" + (i+1), "失败", new ArrayList<>()));
                }
            }

            // 融合所有结果
            String finalResult = consensusMerge(results, transcriptionResults);
            System.out.println("\n🎯 最终融合结果: '" + finalResult + "'");

            // 保存结果到文件
            saveResultToFile(finalResult, audioPath);

            return finalResult;

        } catch (Exception e) {
            throw new RuntimeException("转录失败: " + e.getMessage(), e);
        }
    }

    /**
     * 详细分析音频文件
     */
    private AudioFileInfo analyzeAudioFile(String audioPath) {
        AudioFileInfo info = new AudioFileInfo();
        try {
            File audioFile = new File(audioPath);
            info.fileSize = audioFile.length();

            AudioInputStream audioInputStream = AudioSystem.getAudioInputStream(audioFile);
            AudioFormat format = audioInputStream.getFormat();

            info.sampleRate = format.getSampleRate();
            info.channels = format.getChannels();
            info.sampleSize = format.getSampleSizeInBits();
            info.encoding = format.getEncoding().toString();

            // 计算时长（秒）
            long frames = audioInputStream.getFrameLength();
            info.duration = (double) frames / format.getFrameRate();

            audioInputStream.close();

        } catch (Exception e) {
            System.err.println("分析音频文件失败: " + e.getMessage());
        }
        return info;
    }

    /**
     * 使用单个模型识别音频文件 - 修复JSON解析
     */
    protected TranscriptionResult transcribeWithModel(Model model, String audioPath) {
        Recognizer recognizer = null;
        AudioInputStream audioInputStream = null;

        try {
            // 创建识别器，设置采样率
            recognizer = new Recognizer(model, 16000.0f);

            // 检查音频文件
            File audioFile = new File(audioPath);
            if (!audioFile.exists()) {
                return new TranscriptionResult("", 0.0, model.toString(), "文件不存在", new ArrayList<>());
            }

            // 获取音频输入流
            audioInputStream = AudioSystem.getAudioInputStream(audioFile);
            AudioFormat sourceFormat = audioInputStream.getFormat();

            System.out.println("  源音频格式: " + sourceFormat.getSampleRate() + "Hz, " +
                    sourceFormat.getChannels() + "声道, " +
                    sourceFormat.getSampleSizeInBits() + "bit");

            // 转换为Vosk要求的格式
            AudioFormat targetFormat = new AudioFormat(16000, 16, 1, true, false);
            AudioInputStream convertedAudio = AudioSystem.getAudioInputStream(targetFormat, audioInputStream);

            // 读取音频数据并识别
            byte[] buffer = new byte[4096];
            int bytesRead;
            StringBuilder fullText = new StringBuilder();
            List<String> segments = new ArrayList<>(); // 保存所有识别片段
            double totalConfidence = 0;
            int segmentCount = 0;
            int totalBytesRead = 0;

            System.out.print("  处理进度: ");
            while ((bytesRead = convertedAudio.read(buffer)) >= 0) {
                if (bytesRead > 0) {
                    totalBytesRead += bytesRead;

                    // 显示进度
                    if (totalBytesRead % (4096 * 100) == 0) {
                        System.out.print("#");
                    }

                    if (recognizer.acceptWaveForm(buffer, bytesRead)) {
                        String resultJson = recognizer.getResult();
                        String text = extractTextFromResult(resultJson);
                        double confidence = extractConfidenceFromResult(resultJson);

                        // 放宽过滤条件：只要不是完全空就接受
                        if (text != null && !text.trim().isEmpty() && !text.equals(":")) {
                            System.out.print("[" + text + "]");

                            fullText.append(text).append(" ");
                            segments.add(text); // 保存片段
                            totalConfidence += confidence;
                            segmentCount++;
                        }
                    } else {
                        // 获取部分结果用于调试
                        String partialResult = recognizer.getPartialResult();
                        String partialText = extractPartialTextFromResult(partialResult);
                        if (partialText != null && !partialText.trim().isEmpty() && !partialText.trim().equals(":")) {
                            System.out.print("{" + partialText + "}");
                        } else {
                            System.out.print(".");
                        }
                    }
                }
            }
            System.out.println(); // 换行

            // 处理最终结果
            String finalResultJson = recognizer.getFinalResult();
            String finalText = extractTextFromResult(finalResultJson);
            double finalConfidence = extractConfidenceFromResult(finalResultJson);

            if (finalText != null && !finalText.trim().isEmpty() && !finalText.equals(":")) {
                fullText.append(finalText);
                segments.add(finalText); // 保存最终片段
                totalConfidence += finalConfidence;
                segmentCount++;
            }

            // 计算平均置信度
            double avgConfidence = segmentCount > 0 ? totalConfidence / segmentCount : 0;

            String status = segmentCount > 0 ? "成功识别" : "无识别结果";
            if (totalBytesRead == 0) {
                status = "未读取到音频数据";
            }

            String finalTextResult = fullText.toString().trim();

            // 调试信息
            System.out.println("  识别统计: 片段数=" + segmentCount + ", 总字符数=" + finalTextResult.length());
            if (segmentCount > 0) {
                System.out.println("  所有片段: " + String.join(" | ", segments));
            }

            return new TranscriptionResult(
                    finalTextResult,
                    avgConfidence,
                    model.toString(),
                    status,
                    segments
            );

        } catch (Exception e) {
            System.err.println("模型识别异常: " + e.getMessage());
            e.printStackTrace();
            return new TranscriptionResult("", 0.0, model.toString(), "处理失败: " + e.getMessage(), new ArrayList<>());
        } finally {
            // 清理资源
            if (recognizer != null) {
                recognizer.close();
            }
            if (audioInputStream != null) {
                try {
                    audioInputStream.close();
                } catch (IOException e) {
                    System.err.println("关闭音频流失败: " + e.getMessage());
                }
            }
        }
    }

    /**
     * 从Vosk结果JSON中提取文本 - 修复版本
     */
    private String extractTextFromResult(String resultJson) {
        try {
            if (resultJson == null || resultJson.trim().isEmpty()) {
                return "";
            }

            // 调试：显示原始JSON
            if (resultJson.length() < 200) {
                System.out.println("    完整结果JSON: " + resultJson);
            }

            // 查找 "text" 字段
            int textStart = resultJson.indexOf("\"text\"");
            if (textStart == -1) {
                return "";
            }

            // 找到 "text" 后面的冒号
            int colonIndex = resultJson.indexOf(":", textStart);
            if (colonIndex == -1) {
                return "";
            }

            // 找到冒号后面的第一个引号
            int quoteStart = resultJson.indexOf("\"", colonIndex);
            if (quoteStart == -1) {
                return "";
            }

            // 找到结束引号
            int quoteEnd = resultJson.indexOf("\"", quoteStart + 1);
            if (quoteEnd == -1) {
                return "";
            }

            // 提取引号之间的文本
            String text = resultJson.substring(quoteStart + 1, quoteEnd).trim();

            // 调试：显示提取的文本
            if (!text.isEmpty() && !text.equals(":")) {
                System.out.println("    提取文本: '" + text + "'");
            }

            return text;

        } catch (Exception e) {
            System.err.println("解析文本结果失败: " + resultJson);
            return "";
        }
    }

    /**
     * 从Vosk部分结果JSON中提取文本
     */
    private String extractPartialTextFromResult(String resultJson) {
        try {
            if (resultJson == null || resultJson.trim().isEmpty()) {
                return "";
            }

            // 查找 "partial" 字段
            int partialStart = resultJson.indexOf("\"partial\"");
            if (partialStart == -1) {
                return "";
            }

            // 找到 "partial" 后面的冒号
            int colonIndex = resultJson.indexOf(":", partialStart);
            if (colonIndex == -1) {
                return "";
            }

            // 找到冒号后面的第一个引号
            int quoteStart = resultJson.indexOf("\"", colonIndex);
            if (quoteStart == -1) {
                return "";
            }

            // 找到结束引号
            int quoteEnd = resultJson.indexOf("\"", quoteStart + 1);
            if (quoteEnd == -1) {
                return "";
            }

            // 提取引号之间的文本
            String text = resultJson.substring(quoteStart + 1, quoteEnd).trim();

            return text;

        } catch (Exception e) {
            return "";
        }
    }

    /**
     * 从Vosk结果JSON中提取置信度
     */
    private double extractConfidenceFromResult(String resultJson) {
        try {
            if (resultJson.contains("\"conf\"")) {
                int confStart = resultJson.indexOf("\"conf\"") + 6;
                int confEnd = resultJson.indexOf(",", confStart);
                if (confEnd > confStart) {
                    String confStr = resultJson.substring(confStart, confEnd).trim();
                    return Double.parseDouble(confStr);
                }
            }
        } catch (Exception e) {
            // 忽略置信度解析错误
        }
        return 0.0;
    }

    /**
     * 保存结果到文本文件
     */
    private void saveResultToFile(String result, String audioPath) {
        try {
            // 确保输出目录存在
            File outputFile = new File(outputTextPath);
            File parentDir = outputFile.getParentFile();
            if (parentDir != null && !parentDir.exists()) {
                parentDir.mkdirs();
            }

            // 写入结果
            try (BufferedWriter writer = new BufferedWriter(new FileWriter(outputTextPath, true))) {
                writer.write("=== 转录结果 ===\n");
                writer.write("音频文件: " + new File(audioPath).getName() + "\n");
                writer.write("转录时间: " + new Date() + "\n");
                writer.write("转录结果:\n" + result + "\n");
                writer.write("=".repeat(50) + "\n\n");
            }

            System.out.println("\n💾 转录结果已保存到: " + outputTextPath);

        } catch (IOException e) {
            System.err.println("保存结果到文件失败: " + e.getMessage());
        }
    }

    /**
     * 多模型结果融合算法
     */
    private String consensusMerge(List<String> transcripts, List<TranscriptionResult> results) {
        // 过滤空结果和只有冒号的结果
        List<String> validTranscripts = new ArrayList<>();

        for (int i = 0; i < transcripts.size(); i++) {
            String transcript = transcripts.get(i);
            if (transcript != null && !transcript.trim().isEmpty() &&
                    !transcript.replace(":", "").replace(" ", "").isEmpty()) {
                validTranscripts.add(transcript);
                System.out.println("有效结果 " + (i+1) + ": '" + transcript + "'");
            }
        }

        if (validTranscripts.isEmpty()) {
            System.out.println("所有结果都无效，返回空字符串");
            return "";
        }

        System.out.println("有效结果数量: " + validTranscripts.size());

        if (validTranscripts.size() == 1) {
            return validTranscripts.get(0);
        }

        // 选择最长的结果
        String longestResult = "";
        for (String transcript : validTranscripts) {
            if (transcript.length() > longestResult.length()) {
                longestResult = transcript;
            }
        }

        System.out.println("选择最长结果: '" + longestResult + "' (长度: " + longestResult.length() + ")");

        return longestResult;
    }

    /**
     * 音频文件信息类
     */
    private static class AudioFileInfo {
        long fileSize;
        double duration;
        float sampleRate;
        int channels;
        int sampleSize;
        String encoding;
    }

    /**
     * 转录结果类
     */
    public static class TranscriptionResult {
        public final String text;
        public final double confidence;
        public final String modelName;
        public final String status;
        public final List<String> segments;

        public TranscriptionResult(String text, double confidence, String modelName, String status, List<String> segments) {
            this.text = text;
            this.confidence = confidence;
            this.modelName = modelName;
            this.status = status;
            this.segments = segments;
        }
    }

    /**
     * 关闭资源
     */
    public void close() {
        if (executorService != null) {
            executorService.shutdown();
            try {
                if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                    executorService.shutdownNow();
                }
            } catch (InterruptedException e) {
                executorService.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }

        if (models != null) {
            for (Model model : models) {
                try {
                    model.close();
                } catch (Exception e) {
                    System.err.println("关闭模型失败: " + e.getMessage());
                }
            }
            models.clear();
        }
    }

    /**
     * 使用示例
     */
    public static void main(String[] args) {
        MultiModelTranscriber transcriber = null;
        try {
            // 创建多模型转录器
            transcriber = new MultiModelTranscriber();
            transcriber.setOutputTextPath("D:\\mp4totext\\transcription_results.txt");

            // 转录音频文件
            String audioPath = "D:\\mp4totext\\mp4_testfile\\wav\\古籍修复.wav";

            System.out.println("开始处理音频文件: " + audioPath);
            String result = transcriber.transcribeWithConsensus(audioPath);

            System.out.println("\n" + "=".repeat(60));
            System.out.println("🎉 转录完成！最终结果:");
            System.out.println("=".repeat(60));
            System.out.println(result);
            System.out.println("=".repeat(60));

        } catch (Exception e) {
            System.err.println("转录失败: " + e.getMessage());
            e.printStackTrace();
        } finally {
            if (transcriber != null) {
                transcriber.close();
            }
        }
    }
}